from ultralytics import YOLO
import os
from ultralytics.utils import SETTINGS
from argparse import ArgumentParser
def parse_key_value_pairs(arg_list):
    result = {}
    for item in arg_list:
        key, value = item.split('=')
        try:
            # Try converting to int first
            value = int(value)
        except ValueError:
            try:
                # If conversion to int fails, try converting to float
                value = float(value)
            except ValueError:
                # If both conversions fail, keep the value as a string
                pass
        result[key] = value
    return result
def get_args():
    parser = ArgumentParser()
    parser.add_argument('--weights', type=str ,default='../../weights/yolov8n.pt' ,help='weights path')
    parser.add_argument('--resume', action='store_true', help='resume training')
    parser.add_argument('--epochs', default=20, type=int, help='number of epochs')
    parser.add_argument('--log-dir', default='/project/train/tensorboard/face_det', help='log directory')
    parser.add_argument('--batch', default=-1, type=int, help='size of each image batch')
    parser.add_argument('--imgsz', default=1600, type=int, help='inference size (pixels)')
    parser.add_argument('--train-args', nargs='+')

    # parser.add_argument_group
    # add other params without specify name, for example mosaic=0, dense=1


    return parser.parse_args()
if __name__ == '__main__':
    
    SETTINGS.update(wandb=False)
    print(SETTINGS['wandb'])
    # os.environ['WANDB_DISABLED'] = 'true'
    args = get_args()
    # Extract key-value pairs from the command line
    other_args_dict = parse_key_value_pairs(args.train_args) if args.train_args else {}

    # Combine all arguments and pass them to the train function
    all_args = {**vars(args), **other_args_dict}
    # all_args = { **other_args_dict, **vars(args),}

    # done with args
    print('args: \n', args)
    print('all_args, \n', all_args)
    # Create a new YOLO model from scratch
    # model = YOLO('yolov8m.yaml')

    # Load a pretrained YOLO model (recommended for training)
    if not os.path.isfile(args.weights):
        raise Exception('Invalid weight path')
    model = YOLO(args.weights)
    with open('model.txt', 'w') as f:
        print(model, file=f)
   
    # Train the model using the 'coco128.yaml' dataset for 3 epochs
    if os.getenv('face_dev'):
        cfg = 'facedet_dev.yaml'
    else:
        cfg = 'facedet.yaml'
    print(cfg)
    imgsz= 960
    # imgsz = 640
    # results = model.train(data=cfg, epochs=args.epochs, project='/project/train/models/face_det', resume=args.resume, log_dir = args.log_dir, imgsz=args.imgsz, batch=args.batch_size, **vars(all_args))
    results = model.train(data=cfg, project='/project/train/models/face_det', **(all_args))

    with open('model1.txt', 'w') as f:
        print(model, file=f)
   
    # Evaluate the model's performance on the validation set
    # results = model.val()

    # Perform object detection on an image using the model
    # results = model('https://ultralytics.com/images/bus.jpg')

    # Export the model to ONNX format
    # success = model.export(format='onnx')